package com.icef0x.tools;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.security.SecureClassLoader;
import java.util.Hashtable;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.SimpleJavaFileObject;
import javax.tools.JavaFileObject.Kind;

/**
 * @author ywang
 */
public class ClassFileManager extends ForwardingJavaFileManager {
    
    Hashtable<String, JavaClassObject> classes = new Hashtable<String, JavaClassObject>();
    ClassLoader classloader;
    
    protected ClassFileManager(JavaFileManager fileManager) {
        super(fileManager);
    }
    
    @Override
    public ClassLoader getClassLoader(Location location) {
        if (classloader == null) {
            synchronized (this) {
                if (classloader == null) {
                    classloader = new SecureClassLoader() {
                        protected java.lang.Class<?> findClass(String name) throws ClassNotFoundException {
                            byte[] bytes = getClassBytecode(name);
                            return bytes == null ? super.findClass(name) : super.defineClass(name, bytes, 0,
                                    bytes.length);
                        };
                    };
                }
            }
        }
        return classloader;
    }
    
    public byte[] getClassBytecode(String name) {
        JavaClassObject jClassObject = classes.get(name);
        return jClassObject == null ? null : jClassObject.getBytes();
    }
    
    @Override
    public JavaFileObject getJavaFileForOutput(Location location, String className, Kind kind, FileObject sibling)
            throws IOException {
        JavaClassObject jClassObject = new JavaClassObject(className, kind);
        classes.put(className, jClassObject);
        return jClassObject;
    }
    
    private static class JavaClassObject extends SimpleJavaFileObject {
        
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        
        protected JavaClassObject(String name, Kind kind) {
            super(URI.create("string:///" + name.replace(".", "/") + kind.extension), kind);
        }
        
        public byte[] getBytes() {
            return bos.toByteArray();
        }
        
        @Override
        public OutputStream openOutputStream() throws IOException {
            return bos;
        }
    }
}
